package com.example.nacosprovider.controller;

import org.apache.atlas.AtlasClientV2;
import org.apache.atlas.AtlasServiceException;
import org.apache.atlas.model.discovery.AtlasSearchResult;
import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.model.typedef.AtlasEntityDef;
import org.apache.atlas.model.typedef.AtlasStructDef;
import org.apache.atlas.model.typedef.AtlasTypesDef;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.*;

/**
 * @author ZhangYuan
 * @create 2024/11/13
 * @description
 */
@RestController
public class AtlasController {

    @Autowired
    private AtlasClientV2 atlasClientV2;

    /**
     * 创建实体分类定义
     * @return
     * @throws AtlasServiceException
     */
    @GetMapping("/createType")
    public String createType() throws AtlasServiceException {
        // 创建分类定义
        AtlasTypesDef atlasTypesDef = new AtlasTypesDef();
        // 创建实体分类定义
        AtlasEntityDef atlasEntityDef = new AtlasEntityDef();
        // 创建实体分类属性定义
        AtlasStructDef.AtlasAttributeDef atlasAttributeDef = new AtlasStructDef.AtlasAttributeDef();
        //字段名
        atlasAttributeDef.setName("test_field");
        //字段类型
        atlasAttributeDef.setTypeName("string");
        //是否必填
        atlasAttributeDef.setIsOptional(false);
        ArrayList<AtlasStructDef.AtlasAttributeDef> attributeDefs = new ArrayList<>();
        attributeDefs.add(atlasAttributeDef);
        // 设置实体分类属性
        atlasEntityDef.setAttributeDefs(attributeDefs);
        // 设置实体分类名
        atlasEntityDef.setName("test_entity");
        ArrayList<AtlasEntityDef> entityDefs = new ArrayList<>();
        entityDefs.add(atlasEntityDef);
        atlasTypesDef.setEntityDefs(entityDefs);
        atlasClientV2.createAtlasTypeDefs(atlasTypesDef);
        return "创建成功";
    }

    /**
     * 创建实体
     * @return
     * @throws AtlasServiceException
     */
    @GetMapping("/createEntity")
    public String createEntity() throws AtlasServiceException {
        AtlasEntity atlasEntity = new AtlasEntity();
        atlasEntity.setCreatedBy("admin");
        atlasEntity.setTypeName("test_entity");
        atlasEntity.setCreateTime(new Date());
        atlasEntity.setUpdateTime(new Date());
        atlasEntity.setUpdatedBy("admin");
        Map<String, Object> attributes = new HashMap<>(4);
        attributes.put("test_field", "test_value4");
        atlasEntity.setStatus(AtlasEntity.Status.ACTIVE);
        atlasEntity.setAttributes(attributes);
        AtlasEntity.AtlasEntityWithExtInfo entity = new AtlasEntity.AtlasEntityWithExtInfo();
        entity.setEntity(atlasEntity);
        atlasClientV2.createEntity(entity);
        return "创建成功";
    }

    /**
     * 更新实体
     * @return
     * @throws AtlasServiceException
     */
    @GetMapping("/updateEntity")
    public String updateEntity() throws AtlasServiceException {
        AtlasEntity atlasEntity = new AtlasEntity();
        atlasEntity.setTypeName("test_entity");
        atlasEntity.setUpdateTime(new Date());
        atlasEntity.setUpdatedBy("admin");
        atlasEntity.setGuid("446cc584-6135-4634-ae00-2949611ba11b");
        Map<String, Object> attributes = new HashMap<>(4);
        attributes.put("test_field", "test_value4");
        atlasEntity.setStatus(AtlasEntity.Status.ACTIVE);
        atlasEntity.setAttributes(attributes);
        AtlasEntity.AtlasEntityWithExtInfo entity = new AtlasEntity.AtlasEntityWithExtInfo();
        entity.setEntity(atlasEntity);
//        atlasClientV2.createEntity(entity);//也可执行更新操作
        atlasClientV2.updateEntity(entity);
        return "更新成功";
    }
    /**
     * 查询实体
     * @return
     * @throws AtlasServiceException
     */
    @GetMapping("/getEntity")
    public String getEntity() throws AtlasServiceException {
        AtlasEntity.AtlasEntityWithExtInfo entity = atlasClientV2.getEntityByGuid("446cc584-6135-4634-ae00-2949611ba11b");
        return entity.getEntity().getAttributes().toString();
    }

    /**
     * 删除实体
     * @return
     * @throws AtlasServiceException
     */
    @GetMapping("/deleteEntity")
    public String deleteEntity() throws AtlasServiceException {
        atlasClientV2.deleteEntityByGuid("446cc584-6135-4634-ae00-2949611ba11b");
        return "删除成功";
    }

    /**
     * 基础查询
     * @return
     * @throws AtlasServiceException
     */
    @GetMapping("/query")
    public AtlasSearchResult query() throws AtlasServiceException {
        return atlasClientV2.basicSearch("test_entity", null, "zhangyuan", true, 10, 0);
    }
}
